import math
from tqdm import tqdm
from mydataset import MyDataset
import torch
from torch import nn
from torchvision.transforms import transforms
from cnn_net import CNN_NET
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
import datetime





myTransforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((20, 50))

])


#准备数据集
train_data = MyDataset(r'.\data\train\img',r'.\data\train\label\labels.csv', transform=myTransforms)
test_data = MyDataset(r'.\data\val\img',r'.\data\val\label\labels.csv', transform=myTransforms)


train_data_size=len(train_data)
test_data_size=len(test_data)
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度：{}".format(test_data_size))

def get_batch(data):
    data_size = len(data)
    batch_size=math.floor(data_size*0.1)
    if batch_size<4:
        batch_size=4
    elif batch_size>64:
        batch_size=64
    return batch_size


#利用loader来加载数据集
train_loader=DataLoader(train_data, batch_size=get_batch(train_data),shuffle=True)
test_loader=DataLoader(test_data, batch_size=get_batch(test_data),shuffle=True)

#创建网络模型
net=CNN_NET()

#损失函数
loss_fn= nn.MSELoss()

#优化器
optimizer = torch.optim.Adam(net.parameters(), lr=0.0005)

#设置训练网络的一些参数
#记录训练的次数
total_train_step=0
#记录测试的次数
total_test_step=0
#训练的轮数
epoch=60


start_time=time.process_time()

for i in range(epoch):
    loop = tqdm((train_loader), total=len(train_loader))

    net.train()
    #训练步骤开始
    for data in loop:
        img,label=data
        output=net(img)

        loss=loss_fn(output, label)

        #优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step+=1

        loop.set_description(f'Epoch [{i}/{epoch}]')
        loop.set_postfix(loss=loss.item())




    net.eval()
    #测试步骤开始
    total_test_loss=0
    total_accuracy=0
    #无梯度模式下测试
    with torch.no_grad():
        for data in test_loader:
            imgs, label = data
            output = net(imgs)
            loss = loss_fn(output, label)
            total_test_loss+=loss.item()
            # accuracy=output.eq(label).sum().item()
            # total_accuracy+=accuracy

        # print("整体测试集上的准确率：{}".format(total_accuracy/test_data_size))

        # writer.add_scalar('test_accuracy',total_accuracy/test_data_size,total_test_step)
        total_test_step+=1






#保存模型--最终
torch.save(net,"eyes.pth")
print("模型已保存")


